# 导入所需的库
import cv2, Leap, math, ctypes
import numpy as np
from time import sleep
import cPickle as pickle 

# 定义 Leap Motion 控制器类
class LeapController(object):
    # 初始化函数
    def __init__(self):
        # 创建 Leap Motion 控制器对象
        self.controller = Leap.Controller()
        # 允许访问原始图像数据
        self.controller.set_policy(Leap.Controller.POLICY_IMAGES)
        # 用于存储图像的列表
        self.images = []

    # 拍摄快照函数
    def takeSnapshot(self):
        # 获取 Leap Motion 控制器的当前帧
        frame = self.controller.frame()
        # 仅获取左侧图像
        image = frame.images[0]

        # 将图像数据封装成 NumPy 数组
        i_address = int(image.data_pointer)
        ctype_array_def = ctypes.c_ubyte * image.height * image.width
        # 转换为 ctypes 数组
        as_ctype_array = ctype_array_def.from_address(i_address)
        # 转换为 NumPy 数组
        as_numpy_array = np.ctypeslib.as_array(as_ctype_array)
        rawImage = np.reshape(as_numpy_array, (image.height, image.width))

        # 获取图像的畸变矫正参数
        left_coordinates, left_coefficients = convert_distortion_maps(frame.images[0])

        # 使用 OpenCV 中的 remap 函数进行畸变矫正
        destination = cv2.remap(rawImage, left_coordinates, left_coefficients, interpolation=cv2.INTER_LINEAR)
        # 调整图像大小
        destination = cv2.resize(destination, (400, 400), 0, 0, cv2.INTER_LINEAR)

        # 将处理后的图像添加到图像列表中
        self.images.append(destination)

        print(str(len(self.images)) + ' images captured!')

    # 将图像保存为 pickle 文件函数
    def pickleImages(self, fileName):
        pickleFileName = fileName + ".pickle"
        # 打开 pickle 文件并将图像列表保存到文件中
        pickleFile = open(pickleFileName, 'wb')
        pickle.dump(self.images, pickleFile, pickle.HIGHEST_PROTOCOL)
        pickleFile.close()


# 函数：将 Leap Motion 图像的畸变矫正参数转换为坐标映射和系数映射
def convert_distortion_maps(image):
    distortion_length = image.distortion_width * image.distortion_height
    xmap = np.zeros(distortion_length//2, dtype=np.float32)
    ymap = np.zeros(distortion_length//2, dtype=np.float32)

    for i in range(0, distortion_length, 2):
        xmap[distortion_length//2 - i//2 - 1] = image.distortion[i] * image.width
        ymap[distortion_length//2 - i//2 - 1] = image.distortion[i + 1] * image.height

    xmap = np.reshape(xmap, (image.distortion_height, image.distortion_width//2))
    ymap = np.reshape(ymap, (image.distortion_height, image.distortion_width//2))

    # 调整畸变映射大小以匹配目标图像大小
    resized_xmap = cv2.resize(xmap,
                                (image.width, image.height),
                                0, 0,
                                cv2.INTER_LINEAR)
    resized_ymap = cv2.resize(ymap,
                                (image.width, image.height),
                                0, 0,
                                cv2.INTER_LINEAR)

    # 转换为更快的固定点映射
    coordinate_map, interpolation_coefficients = cv2.convertMaps(resized_xmap,
                                                                    resized_ymap,
                                                                    cv2.CV_32FC1,
                                                                    nninterpolation=False)

    return coordinate_map, interpolation_coefficients

# 函数：将 Leap Motion 图像进行畸变矫正
def undistort(image, coordinate_map, coefficient_map, width, height):
    destination = np.empty((width, height), dtype=np.ubyte)

    # 将图像数据封装成 NumPy 数组
    i_address = int(image.data_pointer)
    ctype_array_def = ctypes.c_ubyte * image.height * image.width
    # 转换为 ctypes 数组
    as_ctype_array = ctype_array_def.from_address(i_address)
    # 转换为 NumPy 数组
    img = np.reshape(as_numpy_array, (image.height, image.width))

    # 使用 remap 函数将图像映射到目标图像
    destination = cv2.remap(img,
                            coordinate_map,
                            coefficient_map,
                            interpolation=cv2.INTER_LINEAR)

    # 调整输出大小以匹配目标图像大小
    destination = cv2.resize(destination,
                            (width, height),
                            0, 0,
                            cv2.INTER_LINEAR)
    return destination


# 函数：对图像进行畸变矫正
def undistortNumpyArray(img, coordinate_map, coefficient_map, width, height):
    destination = np.empty((width, height), dtype=np.ubyte)
    
    # 使用 remap 函数将图像映射到目标图像
    destination = cv2.remap(img,
                            coordinate_map,
                            coefficient_map,
                            interpolation=cv2.INTER_LINEAR)
    
    # 调整输出大小以匹配目标图像大小
    destination = cv2.resize(destination,
                            (width, height),
                            0, 0,
                            cv2.INTER_LINEAR)
    
    return destination

# 函数：将 Leap Motion 图像的畸变矫正参数转换为坐标映射和系数映射
def convertRawImage(img, frame, image):
    maps_initialized = False

    if image.is_valid:
        if not maps_initialized:
            left_coordinates, left_coefficients = convert_distortion_maps(frame.images[0])
            right_coordinates, right_coefficients = convert_distortion_maps(frame.images[1])
            maps_initialized = True

        undistorted_left = undistortNumpyArray(img, left_coordinates, left_coefficients, 400, 400)
        return undistorted_left
    else:
        return 0

# 函数：将 Leap Motion 图像的畸变矫正参数转换为坐标映射和系数映射
def convertImages(frame, image):
    maps_initialized = False

    if image.is_valid:
        if not maps_initialized:
            left_coordinates, left_coefficients = convert_distortion_maps(frame.images[0])
            right_coordinates, right_coefficients = convert_distortion_maps(frame.images[1])
            maps_initialized = True

        undistorted_left = undistort(image, left_coordinates, left_coefficients, 400, 400)
        undistorted_right = undistort(image, right_coordinates, right_coefficients, 400, 400)

        return undistorted_left, undistorted_right
    else:
        return 0

# 函数：拍摄并保存图像
def snapNSave(picCount=0, directory='pictures/'):
    # 创建 Leap Motion 控制器对象
    controller = Leap.Controller()
    # 允许访问原始图像数据
    controller.set_policy(Leap.Controller.POLICY_IMAGES)

    # 等待一段时间以确保控制器初始化完成
    sleep(0.1)

    # 获取当前帧和左侧图像
    frame = controller.frame()
    image = frame.images[0]

    # 将左右图像进行畸变矫正并保存
    [left, right] = convertImages(frame, image)
    cv2.imwrite(directory + 'left' + str(picCount) + '.png', left)
    cv2.imwrite(directory + 'right' + str(picCount) + '.png', right)
    picCount = picCount + 1

    return picCount

# 函数：获取手部模型的相关部分信息
def getModelParts(frame):
    finger_names = ['Thumb', 'Index', 'Middle', 'Ring', 'Pinky']
    bone_names = ['Metacarpal', 'Proximal', 'Intermediate', 'Distal']

    # 获取手部信息
    for hand in frame.hands:
        handType = "Left hand" if hand.is_left else "Right hand"

        print("  %s, id %d, position: %s" % (
            handType, hand.id, hand.palm_position))

        # 获取手的法向量和方向
        normal = hand.palm_normal
        direction = hand.direction

        # 计算手的俯仰、横滚和偏航角度
        print("  pitch: %f degrees, roll: %f degrees, yaw: %f degrees" % (
            direction.pitch * Leap.RAD_TO_DEG,
            normal.roll * Leap.RAD_TO_DEG,
            direction.yaw * Leap.RAD_TO_DEG))

        # 获取手臂骨骼信息
        arm = hand.arm
        print("  Arm direction: %s, wrist position: %s, elbow position: %s" % (
            arm.direction,
            arm.wrist_position,
            arm.elbow_position))

        # 获取手指信息
        for finger in hand.fingers:
            print("    %s finger, id: %d, length: %fmm, width: %fmm" % (
                finger_names[finger.type],
                finger.id,
                finger.length,
                finger.width))

            # 获取手指骨骼信息
            for b in range(0, 4):
                bone = finger.bone(b)
                print("      Bone: %s, start: %s, end: %s, direction: %s" % (
                    bone_names[bone.type],
                    bone.prev_joint,
                    bone.next_joint,
                    bone.direction))

    # 获取工具信息
    for tool in frame.tools:
        print("  Tool id: %d, position: %s, direction: %s" % (
            tool.id, tool.tip_position, tool.direction))


def takeSnapshotWithFinger(numFingers=1, picCount=0, directory='Pickles/'):

    # 创建 Leap Motion 控制器对象
    controller = Leap.Controller()
    # 允许访问原始图像数据
    controller.set_policy(Leap.Controller.POLICY_IMAGES)

    # 等待一段时间以确保控制器初始化完成
    sleep(0.1)

    # 获取当前帧和左侧图像
    frame = controller.frame()
    image = frame.images[0]

    #wrap image data in numpy array
    i_address = int(image.data_pointer)
    ctype_array_def = ctypes.c_ubyte * image.height * image.width
    # as ctypes array
    as_ctype_array = ctype_array_def.from_address(i_address)
    # as numpy array
    as_numpy_array = np.ctypeslib.as_array(as_ctype_array)
    rawImage = np.reshape(as_numpy_array, (image.height, image.width))

    # 获取左右图像的畸变矫正参数
    left_coordinates, left_coefficients = convert_distortion_maps(frame.images[0])
    right_coordinates, right_coefficients = convert_distortion_maps(frame.images[1])

    # 存储手指追踪信息的字典
    trackingIndices = {}

    # 循环遍历手指
    for fingerNum in np.arange(1, numFingers+1):
        # 获取手指顶端和基部的点
        horizontal_slope = -1 * (frame.hands[0].fingers[fingerNum].tip_position.x - 20) / frame.hands[0].fingers[fingerNum].tip_position.y
        vertical_slope = frame.hands[0].fingers[fingerNum].tip_position.z / frame.hands[0].fingers[fingerNum].tip_position.y

        # 将手指顶端点映射到图像坐标
        pixel = image.warp(Leap.Vector(horizontal_slope, vertical_slope, 0))
        pixelIndices = [np.round(pixel.y), np.round(pixel.x)]

        # 创建一个临时数组用于保存手指顶端点
        temp = np.zeros((image.height, image.width))
        temp[pixelIndices[0], pixelIndices[1]] = 255

        # 使用畸变矫正参数将手指顶端点映射到目标图像
        tempMapped = cv2.remap(temp, left_coordinates, left_coefficients, interpolation=cv2.INTER_LINEAR)
        tempMapped = cv2.resize(tempMapped, (400, 400), 0, 0, cv2.INTER_LINEAR)

        # 获取手指顶端点的坐标
        tipIndices = np.where(tempMapped > 1)
        trackingIndices[str(fingerNum)] = {}
        trackingIndices[str(fingerNum)]['tipIndices'] = tipIndices

        # 获取手指另一端的骨头
        bone = frame.hands[0].fingers[fingerNum].bone(2).prev_joint
        horizontal_slope = -1 * (bone.x - 20) / bone.y
        vertical_slope = bone.z / bone.y

        # 将手指基部点映射到图像坐标
        pixel = image.warp(Leap.Vector(horizontal_slope, vertical_slope, 0))
        pixelIndices = [np.round(pixel.y), np.round(pixel.x)]

        # 创建一个临时数组用于保存手指基部点
        temp = np.zeros((image.height, image.width))
        temp[pixelIndices[0], pixelIndices[1]] = 255

        # 使用畸变矫正参数将手指基部点映射到目标图像
        tempMapped = cv2.remap(temp, left_coordinates, left_coefficients, interpolation=cv2.INTER_LINEAR)
        tempMapped = cv2.resize(tempMapped, (400, 400), 0, 0, cv2.INTER_LINEAR)

        # 获取手指基部点的坐标
        baseIndices = np.where(tempMapped > 1)
        trackingIndices[str(fingerNum)]['baseIndices'] = baseIndices

    # 将原始图像进行畸变矫正并保存
    destination = cv2.remap(rawImage, left_coordinates, left_coefficients, interpolation=cv2.INTER_LINEAR)
    destination = cv2.resize(destination, (400, 400), 0, 0, cv2.INTER_LINEAR)

    # 创建图像字典，包含图像、手指追踪信息、图片计数和手指数量
    imageDict = {}
    imageDict['image'] = destination
    imageDict['trackingIndices'] = trackingIndices
    imageDict['picCount'] = picCount
    imageDict['numFingers'] = numFingers

    # 生成 pickle 文件的名称
    pickleName = str(numFingers) + 'f' + str(picCount)
    pickleFileName = directory + pickleName + ".pickle"
    
    # 将图像字典保存为 pickle 文件
    pickleFile = open(pickleFileName, 'wb')
    pickle.dump(imageDict, pickleFile, pickle.HIGHEST_PROTOCOL)
    pickleFile.close()

    return picCount + 1


















